# 打印ncbi的3D json格式的分子数据文件为体素方块
# 需要配合element_data.txt使用
# 只实现了最基本的功能
# 警告：可能有bug，仅供娱乐，不适合学术用途！
# Gitee Repo

import json
import numpy as np
import tkinter
from tkinter import filedialog
import os
import plotly.graph_objects as go
from plotly.offline import plot
import warnings

#0,1,2,3,...,
#0序号元素不被使用
#数据来源：ncbi
element_data_array = [];
if os.path.exists('element_data.txt'):
    element_data_array = np.genfromtxt('element_data.txt',delimiter=',')
else:
    element_data_array = np.zeros((100,3))
    element_data_array[:,0] = np.arange(0,100)
    element_data_array[:,1] = 100
    print('无法打开元素数据文件。无法计算半径和质量')
element_radius_array = element_data_array[:,1]
element_mass_array = element_data_array[:,2]


def get_json_file_path():
    # 创建Tkinter根窗口并立即隐藏，避免出现空白窗口
    root = tkinter.Tk()
    root.withdraw()

    # 打开文件选择对话框，默认过滤为.json文件
    file_path = ''
    file_path = filedialog.askopenfilename(
        title="请选择一个JSON文件",
        filetypes=[("JSON files", "*.json"), ("All files", "*.*")]
    )

    # 检查文件路径是否存在且不是空字符串（即用户选择了文件并且没有取消）
    if file_path and os.path.exists(file_path) and os.path.isfile(file_path):
        print(f"你选择的文件是：{file_path}")
        return file_path
    else:
        raise ValueError("未选择任何文件/文件不存在。")
    return file_path

def get_radius_by_element_id(element_number):
    # element_number = 1,2,3,...
    # unit in pm (10^(-12))
    if element_number <=0 or element_number > element_radius_array.shape[0] - 1:
        raise ValueError('Not Implented')
    return element_radius_array[element_number]/100; #A

def get_mass_by_element_id(element_number):
    # element_number = 1,2,3,...
    # unit in g/mol
    if element_number <=0 or element_number > element_radius_array.shape[0] - 1:
        raise ValueError('Not Implented')
    return element_mass_array[element_number];

def read_json_file(file_path):
    molecule_mass = 0
    atom_coords = [];
    atom_types = [];
    
    with open(file_path, 'r') as file:
        data = json.load(file)
    
    coords_x = np.array( data["PC_Compounds"][0]["coords"][0]["conformers"][0]["x"])
    coords_y = np.array( data["PC_Compounds"][0]["coords"][0]["conformers"][0]["y"])
    coords_z = np.array( data["PC_Compounds"][0]["coords"][0]["conformers"][0]["z"])
    
    atom_coords = np.column_stack((coords_x,coords_y,coords_z))
    atom_types = np.array( data["PC_Compounds"][0]["atoms"]["element"])
    
    if atom_coords.shape[0] != atom_types.shape[0]:
        raise ValueError('数据有误')
    
    molecular_id = int(data["PC_Compounds"][0]["id"]["id"]["cid"])
    molecular_atom_num = atom_coords.shape[0]
        
    for i in range(atom_types.shape[0]):
        molecule_mass += get_mass_by_element_id(atom_types[i])
        
    print('Molecular ID: ' + str(molecular_id))
    print('Molecular Atom Number: ' + str(molecular_atom_num))
    print('Molecular Mass: ' + str(molecule_mass))
    return molecular_atom_num, atom_coords, atom_types

def add_sphere(x,y,z,cx,cy,cz,r):
    cu = np.sqrt((x-cx)**2+(y-cy)**2+(z-cz)**2)*(2/r)
    return np.exp(-cu)*(2/r)

filepath = get_json_file_path()
molecular_atom_num,atom_coords,atom_types = read_json_file(filepath)

# Create Grid
xmax = atom_coords[:,0].max()+2
xmin = atom_coords[:,0].min()-2
ymax = atom_coords[:,1].max()+2
ymin = atom_coords[:,1].min()-2
zmax = atom_coords[:,2].max()+2
zmin = atom_coords[:,2].min()-2

dx = 0.1
x,y,z=np.meshgrid(np.arange(xmin,xmax,dx),np.arange(ymin,ymax,dx),np.arange(zmin,zmax,dx))
x,y,z=x.flatten(),y.flatten(),z.flatten()

# Voxel Data
u = np.zeros_like(x)
r = 1.0

for k in range(molecular_atom_num):
    r = get_radius_by_element_id(atom_types[k]);
    u += add_sphere(x,y,z,atom_coords[k,0],atom_coords[k,1],atom_coords[k,2],r)

fig = go.Figure(data=go.Volume(
    x=x,
    y=y,
    z=z,
    value=u,
    opacity=0.1, # needs to be small to see through all surfaces
    surface_count=20, # needs to be a large number for good volume rendering
    ))
plot(fig)
